package org.usmile.algorithms.learn.tree;

/**
 * 二叉查找树
 */
public class BinarySearchTree<E extends Comparable<E>> {
    private TreeNode<E> root;

    private int size;

    public BinarySearchTree() {
        root = null;
        size = 0;
    }

    public int size() {
        return size;
    }

    public boolean isEmpty() {
        return size == 0;
    }

    public void add(E e) {
        if (root == null) {
            root = new TreeNode<>(e);
            size++;

            return;
        }

        TreeNode<E> curr = root;
        while (true) {
            if (e.compareTo(curr.e) == 0) {
                return;
            }
            if (e.compareTo(curr.e) < 0) {
                if (curr.left == null) {
                    curr.left = new TreeNode<>(e);
                    size++;

                    return;
                } else {
                    curr = curr.left;
                }
            } else {
                if (curr.right == null) {
                    curr.right = new TreeNode<>(e);
                    size++;

                    return;
                } else {
                    curr = curr.right;
                }
            }
        }
    }

    public boolean contains(E e) {
        TreeNode<E> curr = root;
        while (curr != null) {
            if (e.compareTo(curr.e) == 0) {
                return true;
            } else if (e.compareTo(curr.e) < 0) {
                curr = curr.left;
            } else {
                curr = curr.right;
            }
        }

        return false;
    }

    public E getMin() {
        if (root == null) {
            return null;
        }

        TreeNode<E> curr = root;
        while (curr.left != null) {
            curr = curr.left;
        }

        return curr.e;
    }

    public E getMax() {
        if (root == null) {
            return null;
        }

        TreeNode<E> curr = root;
        while (curr.right != null) {
            curr = curr.right;
        }

        return curr.e;
    }

    public E removeMin() {
        if (root == null) {
            return null;
        }

        TreeNode<E> parent = null;
        TreeNode<E> min = root;
        while (min.left != null) {
            parent = min;
            min = min.left;
        }

        if (parent == null) {
            root = root.right;
        } else {
            parent.left = min.right;
            min.right = null;
        }

        size--;

        return min.e;
    }

    public E removeMax() {
        if (root == null) {
            return null;
        }

        TreeNode<E> parent = null;
        TreeNode<E> max = root;
        while (max.right != null) {
            parent = max;
            max = max.right;
        }

        if (parent == null) {
            root = root.left;
        } else {
            parent.right = max.left;
            max.left = null;
        }

        size--;

        return max.e;
    }

    public void remove(E e) {
        if (root == null) return;

        TreeNode<E> curr = root;
        TreeNode<E> parent = null;
        // 找到要删除的节点
        while (curr != null && e.compareTo(curr.e) != 0) {
            parent = curr;
            if (e.compareTo(curr.e) < 0) curr = curr.left;
            else curr = curr.right;
        }
        // 如果没有找到需要删除的元素，则直接返回
        if (curr == null) {
            return;
        }

        if (curr.left != null && curr.right != null) {
            // 1. 找到 curr 的右子树的最小值节点
            TreeNode<E> min = curr.right;
            TreeNode<E> minParent = curr;
            while (min.left != null) {
                minParent = min;
                min = min.left;
            }
            // 2. 覆盖需要删除节点的值为最小值
            curr.e = min.e;

            // 3. 删除最小值节点
            curr = min;
            parent = minParent;
        }

        // 删除节点是叶子节点或者仅有一个子树
        TreeNode<E> child; // 用于存储 需要删除节点 的子节点
        if (curr.left != null) {
            child = curr.left;
            // 注意：删除的是根节点的话，也需要将当前节点的 left 置空
            curr.left = null;
        } else if (curr.right != null) {
            child = curr.right;
            // 注意：删除的是根节点的话，也需要将当前节点的 right 置空
            curr.right = null;
        } else {
            child = null;
        }

        if (parent == null) {
            root = child;
        } else if (curr == parent.left) {
            parent.left = child;
        } else if (curr == parent.right) {
            parent.right = child;
        }

        // bug 修复：需要维护好 size
        size--;
    }


    public static void main(String[] args) {
        BinarySearchTree<Integer> tree = new BinarySearchTree<>();
        tree.add(1);
        tree.add(2);
        tree.add(-1);
        tree.add(-2);
        TreeUtil.print(tree.root);

        System.out.println(tree.contains(1));
        System.out.println(tree.contains(-3));
        System.out.println(tree.getMin());
        System.out.println(tree.getMax());
        System.out.println(tree.removeMin());
        TreeUtil.print(tree.root);
        System.out.println(tree.removeMax());
        TreeUtil.print(tree.root);
    }
}
